"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China
"""
import numpy as np
import commpy as cp
import scipy
from scipy.signal import lfilter


def Electrical_filter(input, transfer_funtion, input_sampling_rate, bandwidth = 40e9, filter_order = 4):
    pulstrain_length = len(input)
    freq = input_sampling_rate / 2 * np.linspace(-1, 1, pulstrain_length)
    if transfer_funtion == 'Gaussian':
        Filter = np.exp(-np.log(2) * (1 * freq / bandwidth) ** (2 * filter_order))

    elif transfer_funtion == 'Rectan':
        a, b = np.where(abs(freq) > bandwidth)
        Filter = np.ones(pulstrain_length, 1)
        Filter[a-1] = 1e-10
    else:
        print('ok')

    data = np.fft.fftshift(np.fft.fft(input)) * Filter
    data = np.fft.ifftshift(data)
    output = np.fft.ifft(data)

    return output


def Electrical_filter_POL(input, Sampling_rate_default, filter_type, BW_electrical, order):

    """电滤波器
       Args:
           input[0]: 输入X偏振信号, numpy类型
           input[1]: 输入Y偏振信号, numpy类型
           Sampling_rate_default: 默认采样率
           filter_type: 滤波器类型
           BW_electrical: 带宽
           order: 阶数

       Returns:
           output[0]：输出X偏振信号
           output[1]：输出Y偏振信号
       """
    Tx_Samples_X_DAC = input[0]
    Tx_Samples_Y_DAC = input[1]

    Tx_Samples_X_DAC = Electrical_filter(Tx_Samples_X_DAC, filter_type, Sampling_rate_default, BW_electrical, order)
    Tx_Samples_Y_DAC = Electrical_filter(Tx_Samples_Y_DAC, filter_type, Sampling_rate_default, BW_electrical, order)

    output = [Tx_Samples_X_DAC, Tx_Samples_Y_DAC]

    return output
